import win32process
import win32api
import win32con
import psutil
import sys
import ctypes
from ctypes import wintypes

# 定义缺少的常量
THREAD_GET_CONTEXT = 0x0008
THREAD_SUSPEND_RESUME = 0x0002
THREAD_QUERY_INFORMATION = 0x0040


# 定义STACKFRAME结构和相关函数
class STACKFRAME(ctypes.Structure):
    _fields_ = [
        ("AddrPC", wintypes.ULONG),
        ("AddrReturn", wintypes.ULONG),
        ("AddrFrame", wintypes.ULONG),
        ("AddrStack", wintypes.ULONG),
        ("FuncTableEntry", ctypes.c_void_p),
        ("Params", wintypes.ULONG * 4),
        ("Reserved", wintypes.ULONG * 2),
        ("Far", wintypes.BOOL),
        ("Virtual", wintypes.BOOL),
        ("Reserved4", wintypes.ULONG),
        ("AddrBStore", wintypes.ULONG),
        ("PrologueByteCount", wintypes.ULONG),
    ]


# 加载DbgHelp库
dbghelp = ctypes.windll.DbgHelp


def print_thread_stack(pid, tid=None):
    """打印指定进程中特定线程或所有线程的堆栈"""
    try:
        # 打开进程（使用更具体的权限）
        process_handle = win32api.OpenProcess(
            win32con.PROCESS_VM_READ | win32con.PROCESS_QUERY_INFORMATION,
            False,
            pid
        )

        # 获取所有线程
        if tid is None:
            threads = []
            for proc in psutil.process_iter(['pid', 'name']):
                if proc.info['pid'] == pid:
                    for thread in proc.threads():
                        threads.append(thread.id)
            print(f"进程 {pid} 共有 {len(threads)} 个线程")
        else:
            threads = [tid]

        # 初始化符号处理
        dbghelp.SymInitialize(process_handle, None, True)

        # 打印每个线程的堆栈
        for thread_id in threads:
            try:
                # 打开线程（使用更具体的权限）
                thread_handle = win32api.OpenThread(
                    THREAD_GET_CONTEXT | THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION,
                    False,
                    thread_id
                )

                # 暂停线程以获取稳定的上下文
                win32process.SuspendThread(thread_handle)

                try:
                    # 获取线程上下文
                    context = win32process.GetThreadContext(
                        thread_handle,
                        win32con.CONTEXT_FULL
                    )

                    # 打印线程ID和上下文信息
                    print(f"\n线程ID: {thread_id}")
                    print(f"EIP: 0x{context['Eip']:X}")
                    print(f"ESP: 0x{context['Esp']:X}")

                    # 打印堆栈跟踪
                    print("堆栈跟踪:")
                    stack = STACKFRAME()
                    stack.AddrPC.QuadPart = context['Eip']
                    stack.AddrFrame.QuadPart = context['Ebp']
                    stack.AddrStack.QuadPart = context['Esp']

                    # 展开堆栈
                    frame_number = 0
                    while dbghelp.StackWalk64(
                            win32con.IMAGE_FILE_MACHINE_I386,  # 假设x86架构
                            process_handle,
                            thread_handle,
                            ctypes.byref(stack),
                            ctypes.byref(context),
                            None,
                            dbghelp.SymFunctionTableAccess64,
                            dbghelp.SymGetModuleBase64,
                            None
                    ):
                        # 获取函数名
                        symbol_info = ctypes.create_string_buffer(256)
                        displacement = wintypes.ULONG()
                        if dbghelp.SymGetSymFromAddr64(
                                process_handle,
                                stack.AddrPC.QuadPart,
                                ctypes.byref(displacement),
                                symbol_info
                        ):
                            symbol_name = symbol_info.raw.decode('utf-8', errors='ignore').split('\x00')[0]
                            print(f"#{frame_number}: 0x{stack.AddrPC.QuadPart:X} - {symbol_name}")
                        else:
                            print(f"#{frame_number}: 0x{stack.AddrPC.QuadPart:X} - (unknown)")

                        frame_number += 1
                        if frame_number > 20:  # 限制显示帧数，避免无限循环
                            break

                finally:
                    # 恢复线程执行
                    win32process.ResumeThread(thread_handle)

                win32api.CloseHandle(thread_handle)
            except Exception as e:
                print(f"获取线程 {thread_id} 堆栈失败: {e}")

        # 清理符号处理
        dbghelp.SymCleanup(process_handle)
        win32api.CloseHandle(process_handle)
    except Exception as e:
        print(f"操作失败: {e}")
if __name__ == "__main__":
    pid = 15276
    tid = None

    print_thread_stack(pid, tid)